"""
网络结构
- backbone
  - 卷积层 CBL
    - Conv
    - BN
    - LeakReLu
  - 残差单元 ResUnit
  - 下采样 DownSample
- neck
  - 卷积集合 ConvolutionSet
  - 上采样 UpSample
  - 拼接操作 torch.cat

"""
import torch
from torch import nn


class CBL(nn.Module):
    # Conv+BN+LeakReLu
    def __init__(self, c_in, c_out, k, s):
        super().__init__()
        self.cnn_layer = nn.Sequential(
            nn.Conv2d(c_in, c_out, kernel_size=k, stride=s, padding=k // 2, bias=False),
            nn.BatchNorm2d(c_out),
            nn.LeakyReLU()
        )

    def forward(self, x):
        return self.cnn_layer(x)


class ResUnit(nn.Module):
    # 残差单元
    def __init__(self, c_num):
        super().__init__()
        self.block = nn.Sequential(
            CBL(c_num, c_num // 2, 1, 1),
            CBL(c_num // 2, c_num, 3, 1)
        )

    def forward(self, x):
        return self.block(x) + x


class DownSample(nn.Module):
    # 下采样
    def __init__(self, c_in, c_out):
        super().__init__()
        self.down_sample = nn.Sequential(
            CBL(c_in, c_out, 3, 2)
        )

    def forward(self, x):
        return self.down_sample(x)


class ConvolutionSet(nn.Module):
    # 卷积集合
    def __init__(self, c_in, c_out):
        super().__init__()
        self.cnn_set = nn.Sequential(
            CBL(c_in, c_out, 1, 1),
            CBL(c_out, c_in, 3, 1),
            CBL(c_in, c_out, 1, 1),
            CBL(c_out, c_in, 3, 1),
            CBL(c_in, c_out, 1, 1)
        )

    def forward(self, x):
        return self.cnn_set(x)


class UpSample(nn.Module):
    # 上采样
    def __init__(self):
        super().__init__()
        self.up_sample = nn.Upsample(scale_factor=2, mode='nearest')

    def forward(self, x):
        return self.up_sample(x)


if __name__ == '__main__':
    # data = torch.randn(1, 3, 416, 416)
    # cnn = nn.Sequential(
    #     CBL(3, 32, 3, 1),
    #     DownSample(32, 64)
    # )
    # res = ResUnit(64)
    #
    # cnn_out = cnn(data)
    # res_out = res(cnn_out)
    # print(cnn_out.shape)
    # # torch.Size([1, 64, 208, 208])
    # print(res_out.shape)
    # # torch.Size([1, 64, 208, 208])

    data = torch.randn(1, 1024, 13, 13)
    con_set = ConvolutionSet(1024, 512)
    cnn = CBL(512, 256, 1, 1)
    up_sample = UpSample()

    P0_out = up_sample(cnn(con_set(data)))
    print(P0_out.shape)
    # torch.Size([1, 256, 26, 26])
    pass
